{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.append(\"../..\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from prettytable import PrettyTable, ALL\n",
    "from textwrap import fill\n",
    "from profiler.advisor.interface.interface import Interface"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "profiling_path = \"YOUR PROFILING PATH\"\n",
    "interface = Interface(profiling_path=profiling_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 融合算子API识别"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "指定profiling路径后，可以自动识别其中包含的融合算子并给出对应的torch_npu api和需要修改的代码堆栈。基于给定堆栈可以快速定位到需要修改的代码段，替换torch_npu api后，能够减少pytorch侧的小算子的下发，进而提升模型训练速度。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                                    \r"
     ]
    }
   ],
   "source": [
    "timeline_fusion_ops_result = interface.get_result(\"schedule\", \"timeline_fusion_ops\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<table>\n",
       "    <thead>\n",
       "        <tr>\n",
       "            <th>problem</th>\n",
       "            <th>description</th>\n",
       "            <th>suggestion</th>\n",
       "        </tr>\n",
       "    </thead>\n",
       "    <tbody>\n",
       "        <tr>\n",
       "            <td>timeline_fusion_ops</td>\n",
       "            <td>Found 2 apis to be replaced based on the runtime env cann-8.0.RC1 and torch-2.1.0</td>\n",
       "            <td>1. Please replace training api according to sub table &#x27;Affinity training api&#x27;</td>\n",
       "        </tr>\n",
       "    </tbody>\n",
       "</table>"
      ],
      "text/plain": [
       "+---------------------+---------------------------------------------------------------------------------+-------------------------------------------------------------------------------+\n",
       "|       problem       |                                   description                                   |                                   suggestion                                  |\n",
       "+---------------------+---------------------------------------------------------------------------------+-------------------------------------------------------------------------------+\n",
       "| timeline_fusion_ops | Found 2 apis to be replaced based on the runtime env cann-8.0.RC1 and torch-2.1.0 | 1. Please replace training api according to sub table 'Affinity training api' |\n",
       "+---------------------+---------------------------------------------------------------------------------+-------------------------------------------------------------------------------+"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "display_column_num = 3\n",
    "problems = timeline_fusion_ops_result.get(\"problems\")\n",
    "problem_table = PrettyTable(problems.get(\"headers\")[:display_column_num])\n",
    "for row in problems.get(\"data\"):\n",
    "    for i in range(len(row)):\n",
    "        row[i] = fill(str(row[i]), width=80)\n",
    "    problem_table.add_row(row[:display_column_num])\n",
    "\n",
    "display(problem_table)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "如下所示，存在亲和优化器和梯度裁剪两个可替换的torch_npu api，并给出了具体的堆栈。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<table>\n",
       "    <thead>\n",
       "        <tr>\n",
       "            <th>Affinity API</th>\n",
       "            <th>Code stacks</th>\n",
       "            <th>Stack called counts</th>\n",
       "        </tr>\n",
       "    </thead>\n",
       "    <tbody>\n",
       "        <tr>\n",
       "            <td>optimizer.clip_grad_norm_fused_</td>\n",
       "            <td>/home/ma-user/anaconda3/envs/PyTorch-1.11.0/lib/python3.9/site-<br>packages/torch/nn/utils/clip_grad.py(49): clip_grad_norm_;  /home/ma-<br>user/work/algorithms/doc_cls/Bert.py(205): train_epoch;  /home/ma-<br>user/work/algorithms/doc_cls/Bert.py(252): &lt;module&gt;</td>\n",
       "            <td>2</td>\n",
       "        </tr>\n",
       "        <tr>\n",
       "            <td>torch_npu.optim.NpuFusedAdamW</td>\n",
       "            <td>/home/ma-user/anaconda3/envs/PyTorch-1.11.0/lib/python3.9/site-<br>packages/torch_npu/npu/profiler.py(675): __enter__;  /home/ma-<br>user/anaconda3/envs/PyTorch-1.11.0/lib/python3.9/site-<br>packages/torch_npu/npu/profiler.py(719): wrapper;  /home/ma-<br>user/anaconda3/envs/PyTorch-1.11.0/lib/python3.9/site-<br>packages/torch/optim/lr_scheduler.py(65): wrapper;  /home/ma-<br>user/work/algorithms/doc_cls/Bert.py(219): train_epoch;  /home/ma-<br>user/work/algorithms/doc_cls/Bert.py(252): &lt;module&gt;</td>\n",
       "            <td>2</td>\n",
       "        </tr>\n",
       "    </tbody>\n",
       "</table>"
      ],
      "text/plain": [
       "+---------------------------------+-----------------------------------------------------------------------+---------------------+\n",
       "|           Affinity API          |                              Code stacks                              | Stack called counts |\n",
       "+---------------------------------+-----------------------------------------------------------------------+---------------------+\n",
       "| optimizer.clip_grad_norm_fused_ |    /home/ma-user/anaconda3/envs/PyTorch-1.11.0/lib/python3.9/site-    |          2          |\n",
       "|                                 | packages/torch/nn/utils/clip_grad.py(49): clip_grad_norm_;  /home/ma- |                     |\n",
       "|                                 |   user/work/algorithms/doc_cls/Bert.py(205): train_epoch;  /home/ma-  |                     |\n",
       "|                                 |          user/work/algorithms/doc_cls/Bert.py(252): <module>          |                     |\n",
       "+---------------------------------+-----------------------------------------------------------------------+---------------------+\n",
       "|  torch_npu.optim.NpuFusedAdamW  |    /home/ma-user/anaconda3/envs/PyTorch-1.11.0/lib/python3.9/site-    |          2          |\n",
       "|                                 |     packages/torch_npu/npu/profiler.py(675): __enter__;  /home/ma-    |                     |\n",
       "|                                 |         user/anaconda3/envs/PyTorch-1.11.0/lib/python3.9/site-        |                     |\n",
       "|                                 |      packages/torch_npu/npu/profiler.py(719): wrapper;  /home/ma-     |                     |\n",
       "|                                 |         user/anaconda3/envs/PyTorch-1.11.0/lib/python3.9/site-        |                     |\n",
       "|                                 |     packages/torch/optim/lr_scheduler.py(65): wrapper;  /home/ma-     |                     |\n",
       "|                                 |   user/work/algorithms/doc_cls/Bert.py(219): train_epoch;  /home/ma-  |                     |\n",
       "|                                 |          user/work/algorithms/doc_cls/Bert.py(252): <module>          |                     |\n",
       "+---------------------------------+-----------------------------------------------------------------------+---------------------+"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fusion_ops_api = timeline_fusion_ops_result.get(\"timeline_fusion_ops\")\n",
    "if fusion_ops_api:\n",
    "    fusion_ops_api_table = PrettyTable(fusion_ops_api.get(\"headers\"))\n",
    "\n",
    "    for row in fusion_ops_api.get(\"data\"):\n",
    "        for i in range(len(row)):\n",
    "            row[i] = fill(str(row[i]), width=80)\n",
    "        fusion_ops_api_table.add_row(row)\n",
    "\n",
    "    fusion_ops_api_table.hrules = ALL\n",
    "    display(fusion_ops_api_table)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
